package com.xd.springaiagent.agent;

import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage.ToolResponse;

import java.util.*;

@Slf4j
@EqualsAndHashCode(callSuper = true)
@Component
public class ToolCallAgent extends ReActAgent {
    @Autowired
    protected ToolCallbackProvider toolCallbackProvider;
    protected List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();

    @Override
    protected boolean think() {
        ChatOptions options = ToolCallingChatOptions.builder()
                .toolCallbacks(toolCallbackProvider.getToolCallbacks())
                .internalToolExecutionEnabled(false)
                .build();
        ChatResponse chatResponse = chatClient.prompt(systemPrompt)
                .system(nextStepPrompt)
                .messages(msg)
                .options(options)
                .call()
                .chatResponse();

        AssistantMessage output = chatResponse.getResult().getOutput();
        toolCalls.addAll(output.getToolCalls());
        log.info("==> 本轮挑选的工具：{}", toolCalls);
        String content = output.getText();
        addMessage("assistant", content);

        // 检查是否为终止请求
        if (toolCalls.stream().anyMatch(toolCall -> "terminate".equals(toolCall.name()))) {
            System.out.println("检测到终止请求，结束交互");
            setState(AgentState.FINISHED);
            toolCalls.clear();
            return false;
        }

        return !toolCalls.isEmpty();
    }

    @Override
    protected String act() {
        List<String> results = new ArrayList<>();
        int maxRetries = 3; // 最大重试次数
        int retryCount = 0;

        if (toolCalls.isEmpty()) {
            setState(AgentState.FINISHED);
            return "任务完成";
        }

        for (AssistantMessage.ToolCall toolCall : toolCalls) {
            String toolName = toolCall.name();
            String arguments = toolCall.arguments();
            String toolCallId = toolCall.id();
            String result = null;

            // 重试机制
            while (retryCount < maxRetries) {
                try {
                    ToolCallback toolCallback = Arrays.stream(toolCallbackProvider.getToolCallbacks())
                            .filter(callback -> callback.getToolDefinition().name().equals(toolName))
                            .findFirst()
                            .orElse(null);

                    if (toolCallback == null) {
                        result = "[Error: 未知工具: " + toolName + "]";
                        log.error("工具 {} 不存在，提示 AI 重新选择工具", toolName);
                        addMessage("assistant", "工具 " + toolName + " 不存在，请选择其他有效工具。");
                        if (think()) {
                            // 如果 think() 返回 true，说明生成了新的工具调用，跳出重试
                            break;
                        } else {
                            setState(AgentState.FINISHED);
                            results.add(result);
                            break;
                        }
                    }

                    // 执行工具调用
                    log.info("调用工具: {}，参数: {}", toolName, arguments);
                    result = toolCallback.call(arguments);
                    results.add(result);
                    log.info("工具 {} 执行结果: {}", toolName, result);

                    // 添加工具响应到消息历史
                    ToolResponse toolResponse = new ToolResponse(toolCallId, arguments, result);
                    ToolResponseMessage toolResponseMessage = new ToolResponseMessage(Collections.singletonList(toolResponse));
                    msg.add(toolResponseMessage);

                    break; // 成功执行，跳出重试循环
                } catch (Exception e) {
                    retryCount++;
                    result = "[Error: 工具 " + toolName + " 执行失败: " + e.getMessage() + "]";
                    log.error("工具 {} 执行失败，重试 {}/{}", toolName, retryCount, maxRetries, e);
                    if (retryCount >= maxRetries) {
                        addMessage("assistant", "工具 " + toolName + " 多次尝试失败，请重新选择工具或终止任务。");
                        results.add(result);
                        if (!think()) {
                            setState(AgentState.FINISHED);
                            break;
                        }
                    }
                }
            }
        }
        toolCalls.clear(); // 清空工具调用

        return String.join("\n", results); // 返回结果
    }

    @Override
    protected void cleanup() {
        setState(AgentState.IDLE);
        currentStep = 0;
//        toolCalls.clear();
        // 可选择清空会话历史：
        msg.clear();
    }
}